#!/usr/bin/env bash
# DNS 工具：设置/恢复/查看/对比/测速/一键设置+更新
set -Eeuo pipefail

type log_info >/dev/null 2>&1 || { echo "[dns_tools] _common.sh 未加载"; exit 1; }

# 常用公共 DNS 预设
DNS_PRESETS() {
  cat <<'EOF'
1) AliDNS        : 223.5.5.5 223.6.6.6
2) Google        : 8.8.8.8 8.8.4.4
3) Cloudflare    : 1.1.1.1 1.0.0.1
4) 114DNS        : 114.114.114.114 114.114.115.115
5) Quad9         : 9.9.9.9 149.112.112.112
6) 自定义        : <手动输入>
EOF
}

# 应用 DNS（systemd-resolved 优先）
_apply_dns_servers() {
  ensure_root
  local servers="$1"
  local conf="/etc/systemd/resolved.conf"
  if systemctl list-unit-files | grep -q '^systemd-resolved\.service'; then
    [[ -f "$conf" ]] && cp "$conf" "${conf}.bak.$(date +%F-%T)"
    echo -e "[Resolve]\nDNS=${servers}\nFallbackDNS=" > "$conf"
    ln -sf /run/systemd/resolve/stub-resolv.conf /etc/resolv.conf
    systemctl restart systemd-resolved
    log_info "已通过 systemd-resolved 设置 DNS: ${servers}"
  else
    [[ -f /etc/resolv.conf ]] && cp /etc/resolv.conf "/etc/resolv.conf.bak.$(date +%F-%T)"
    : > /etc/resolv.conf
    for ip in $servers; do echo "nameserver $ip" >> /etc/resolv.conf; done
    log_info "已写入 /etc/resolv.conf: ${servers}"
  fi
}

dns_set_interactive() {
  echo "====== 选择 DNS 方案 ======"
  DNS_PRESETS
  read -rp "选择 [1-6]: " opt
  local servers=""
  case "$opt" in
    1) servers="223.5.5.5 223.6.6.6" ;;
    2) servers="8.8.8.8 8.8.4.4" ;;
    3) servers="1.1.1.1 1.0.0.1" ;;
    4) servers="114.114.114.114 114.114.115.115" ;;
    5) servers="9.9.9.9 149.112.112.112" ;;
    6) read -rp "请输入 DNS 列表: " servers ;;
    *) echo "无效选择"; return 1 ;;
  esac
  [[ -z "$servers" ]] && { echo "未输入 DNS"; return 1; }
  _apply_dns_servers "$servers"
  dns_show_config
}

dns_restore_default() {
  ensure_root
  local conf="/etc/systemd/resolved.conf"
  if systemctl list-unit-files | grep -q '^systemd-resolved\.service'; then
    sed -i -E 's/^\s*DNS=.*$/# DNS=/' "$conf" 2>/dev/null || true
    sed -i -E 's/^\s*FallbackDNS=.*$/# FallbackDNS=/' "$conf" 2>/dev/null || true
    systemctl restart systemd-resolved
    log_info "已恢复 systemd-resolved 默认 DNS"
  else
    local latest="$(ls -t /etc/resolv.conf.bak.* 2>/dev/null | head -n1 || true)"
    [[ -n "$latest" ]] && cp -f "$latest" /etc/resolv.conf || echo -e "nameserver 1.1.1.1\nnameserver 8.8.8.8" > /etc/resolv.conf
    log_info "已恢复 /etc/resolv.conf"
  fi
  dns_show_config
}

dns_show_config() {
  echo "====== 当前 DNS 配置 ======"
  systemctl is-active systemd-resolved >/dev/null 2>&1 && resolvectl status 2>/dev/null || true
  echo "--- /etc/resolv.conf ---"
  cat /etc/resolv.conf 2>/dev/null || true
}

_ensure_dig() {
  command -v dig >/dev/null 2>&1 || { ensure_root; apt_quiet_update; apt-get install -y dnsutils; }
}

dns_compare_resolution() {
  _ensure_dig
  read -rp "输入域名: " domain
  [[ -z "$domain" ]] && return 1
  local servers=("223.5.5.5" "8.8.8.8" "1.1.1.1" "114.114.114.114")
  for s in "${servers[@]}"; do
    echo "---- @$s ----"
    dig +short "@$s" "$domain" A
    dig +short "@$s" "$domain" AAAA
  done
}

dns_pollution_check() {
  _ensure_dig
  read -rp "输入域名: " domain
  [[ -z "$domain" ]] && return 1
  local a1="$(dig +short "@8.8.8.8" "$domain" A | sort -u | tr '\n' ' ')"
  local a2="$(dig +short "@223.5.5.5" "$domain" A | sort -u | tr '\n' ' ')"
  echo "Google: $a1"
  echo "AliDNS: $a2"
  [[ "$a1" != "$a2" ]] && echo "⚠ 结果不同，可能存在污染/劫持" || echo "✓ 解析一致"
}

dns_latency_benchmark() {
  _ensure_dig
  local test_domain="cloudflare.com"
  local servers=("223.5.5.5" "8.8.8.8" "1.1.1.1" "114.114.114.114" "9.9.9.9")
  echo "对 $test_domain 进行 DNS 延迟测试"
  for s in "${servers[@]}"; do
    local t="$(dig +stats "@$s" "$test_domain" | awk '/Query time:/{print $(NF-1)}' | tail -n1)"
    echo "$s : ${t:-"-"} ms"
  done
}

dns_set_and_full_update() {
  dns_set_interactive
  echo ">>> 执行系统完整更新 ..."
  system_update
}
